import math
import numpy as np
import pandas as pd

g = 9.8                       # 重力加速度
missile_speed = 300.0         # 导弹速度 m/s
Rc = 10.0                # 云团半径 m
vc = 3.0              # 下沉速度 m/s（A层仅用于信息展示）
tc = 20.0            # 有效期 s（A层不强制使用）
r = 7.0           # 圆柱半径
h = 10.0          # 圆柱高度

R_tgt_eff = math.sqrt(r**2 + (h/2.0)**2)  # 圆柱斜半径，用于角半径近似

O_fake = np.array([0.0, 0.0, 0.0])           # 假目标
O_target = np.array([0.0, 200.0, 5.0])       # 圆柱轴心中点

pM0 = np.array([20000.0, 0.0, 2000.0])
FY = {
    1: np.array([17800.0,     0.0, 1800.0]),
    2: np.array([12000.0,  1400.0, 1400.0]),
    3: np.array([ 6000.0, -3000.0,  700.0]),
}

V_MIN, V_MAX = 70.0, 140.0

DT = 0.5 # A层搜索网格（可调）
T_hit = np.linalg.norm(O_fake - pM0) / missile_speed # 计算导弹击中假目标时刻
t_grid = np.arange(0.0, min(70.0, T_hit - 1e-3), DT) # 构建搜索网格

DELTA_Z_MAX = 1200.0    # 允许低于机体的最大起爆高度差
NZ = 121                # z_E 采样数（越大越细）
TOP_K = 30              # 每机保留前K个候选

# 打分权重（可微调）
W_dist_over_T = 0.05     # 距离/时间（等于所需平均速，越小越好）
W_tfuse = 0.1           # 引信延时惩罚
W_theta = 0.2           # 与目标方向夹角惩罚
W_z_closeness = 0.15    # 起爆高度贴近机体的奖励
# 角覆盖近似得分（S_tilde）默认权重为 1

def unit(v): # 获取v的单位向量
    n = np.linalg.norm(v)
    if n < 1e-12:
        return v
    return v / n
# 导弹轨迹：直线朝假目标 O_fake 匀速
vM_dir = unit(O_fake - pM0)
vM = missile_speed * vM_dir

def pM(t: float) -> np.ndarray: # 导弹位置
    return pM0 + vM * t

# 小角度近似：球面“角帽”→ 切平面圆盘
# 返回：目标角域被覆盖的近似比例 in [0,1]
# alpha = 云团角半径（弧度），beta = 目标角半径（弧度），theta = 两中心夹角（弧度）
def overlap(alpha: float, beta: float, theta: float) -> float:
    alpha = max(0.0, float(alpha))
    beta = max(1e-9, float(beta))
    theta = max(0.0, float(theta))
    if theta >= alpha + beta:  # 两圆不相交
        return 0.0
    if theta <= abs(beta - alpha): # 目标圆完全包含于云团圆
        if alpha >= beta:
            return 1.0
        else:
            return (alpha * alpha) / (beta * beta)

    # 两圆相交
    r1, r2, d = alpha, beta, theta
    try:
        term1 = r1*r1 * math.acos((d*d + r1*r1 - r2*r2) / (2*d*r1))
        term2 = r2*r2 * math.acos((d*d + r2*r2 - r1*r1) / (2*d*r2))
        term3 = 0.5 * math.sqrt(
            max(0.0, (-d + r1 + r2) * (d + r1 - r2) * (d - r1 + r2) * (d + r1 + r2))
        )

        area = term1 + term2 - term3
        frac = area / (math.pi * r2 * r2)
        return float(np.clip(frac, 0.0, 1.0))
    except ValueError:
        # 数值误差fallback
        return 0.0

# 单机：生成候选起爆位置等信息
# 返回列表：每项 (score, Ex, Ey, heading_x, heading_y, vreq, TE, zE, tfuse, Trel, alpha, beta, theta)
def plan_traj(P0_xy: np.ndarray,
                                  z0: float,
                                  t_grid: np.ndarray,
                                  delta_z_max: float = DELTA_Z_MAX,
                                  nz: int = NZ,
                                  vmin: float = V_MIN,
                                  vmax: float = V_MAX,
                                  top_k: int = TOP_K):
    C = []
    # zE 在 [max(0, z0 - delta), z0]
    z_min = max(0.0, z0 - delta_z_max)
    z_grid = np.linspace(z_min, z0, nz)

    for t in t_grid:
        PM = pM(t)
        d = unit(O_target - PM)  # 导弹到真目标连线
        if abs(d[2]) < 1e-6:
            continue  # 与z平面近乎平行，跳过
        tdir = unit(O_target - PM) # 目标方向单位向量（用于theta）
        for zE in z_grid:
            s = (zE - PM[2]) / d[2] # 与 z=zE 平面的交点
            if s <= 0:
                continue  # 交点在导弹后方或重合，跳过

            E = PM + s * d  # 起爆点（A层把起爆时刻取为 t）
            Exy = E[:2]
            dist = np.linalg.norm(Exy - P0_xy) # 起爆点xy到起飞点的距离
            TE = t
            if TE <= 1e-9:
                continue

            vreq = dist / TE # 要求速度
            if not (vmin <= vreq <= vmax):
                continue

            # 自由落体引信延时（仅竖直方向）
            dz = z0 - zE
            if dz < 0:
                continue  # 起爆高度高于机体，无自由落体解
            tfuse = math.sqrt(max(0.0, 2.0 * dz / g))
            Trel = TE - tfuse
            if Trel < 0:
                continue  # 来不及投放

            # 角域参数（在起爆当刻 t=TE 处的快速近似）
            rM_E = np.linalg.norm(E - PM)
            # 云团角半径 alpha = arcsin(R / r)
            alpha = math.asin(min(1.0, Rc / max(Rc, rM_E)))
            vdir = unit(E - PM)
            theta = math.acos(np.clip(np.dot(vdir, tdir), -1.0, 1.0))
            # 目标角半径（近似）：beta = arctan(R_eff / dMt)
            dMt = np.linalg.norm(O_target - PM)
            beta = math.atan(R_tgt_eff / max(1e-9, dMt))
            S_tilde = overlap(alpha, beta, theta)

            # 打分：角覆盖主导；惩罚大速度需求、长引信、偏离目标方向；奖励 zE 靠近 z0
            score = (S_tilde
                      # - W_dist_over_T * (dist / TE)
                     - W_tfuse * tfuse
                     - W_theta * theta
                      + W_z_closeness * (1.0 - abs(z0 - zE) / max(1.0, delta_z_max)))

            h = (Exy - P0_xy) / max(1e-9, dist)

            C.append((
                float(score), float(Exy[0]), float(Exy[1]),
                float(h[0]), float(h[1]),
                float(vreq), float(TE), float(zE),
                float(tfuse), float(Trel),
                float(alpha), float(beta), float(theta)
            ))

    # 取前K个
    C.sort(key=lambda x: -x[0])
    return C[:top_k]

def main():
    rows = []
    for i in (1, 2, 3):
        P0 = FY[i]
        P0_xy = P0[:2]
        z0 = float(P0[2])

        Points = plan_traj(
            P0_xy=P0_xy,
            z0=z0,
            t_grid=t_grid,
            delta_z_max=DELTA_Z_MAX,
            nz=NZ,
            vmin=V_MIN,
            vmax=V_MAX,
            top_k=TOP_K
        )

        for (score, Ex, Ey, hx, hy, vreq, TE, zE, tfuse, Trel, alpha, beta, theta) in Points:
            heading_deg = math.degrees(math.atan2(hy, hx))
            rows.append({
                "drone": f"FY{i}",
                "score": score,
                "Ex": Ex, "Ey": Ey, "zE": zE,
                "TE": TE, "T_release": Trel, "t_fuse": tfuse,
                "heading_x": hx, "heading_y": hy, "heading_deg": heading_deg,
                "v_required": vreq,
                "alpha(rad)": alpha, "beta(rad)": beta, "theta(rad)": theta
            })

    df = pd.DataFrame(rows).sort_values(["drone", "score"], ascending=[True, False])
    print(df.head(10))  # 预览

    # 导出
    df.to_csv("selected_direct.csv", index=False)

if __name__ == "__main__":
    main()
